#include <linux/kvm_host.h>
#include <linux/types.h>
#include <linux/kvm_host.h>
#include <linux/kernel.h>
#include <linux/export.h>
#include <linux/in.h>
#include <asm/types.h>
#include "kvm_cache_regs.h"
#include "x86.h"
#include "mmu.h"
#include "kvm_vmi.h"
#include "crumbs.h"

int start_trace_from_host(struct kvm * kvm, char * tag)
{
	int i;
	//u16 j;
	//struct kvm_sregs sregs;
	//u8 *idt;
	//u64 idt_base,efer;
	//u32 error;
	//unsigned long cr4,cr0;
	struct kvm_vcpu *vcpu;
	struct msr_data msr;

	if(kvm->crumbs_data.running==1){
		printk("kvm:start_syscall_trace: WARNING: crumbs is already running, start will be aborted.\n");
		return 1;
	}

	vcpu=kvm_get_vcpu(kvm,0);

	if(!is_protmode(vcpu)){
		printk("kvm:start_syscall_trace: ERROR: guest is running in real mode, nitro can not function.\n");
		return 3;
	}

	kvm->crumbs_data.mode = PROT;
	kvm->crumbs_data.idt_entry_size=8;

	if(is_pae(vcpu)){
		kvm->crumbs_data.mode = PAE;
		printk("kvm:start_syscall_trace: system running in PAE mode\n");
	}
	if(is_long_mode(vcpu)){
		kvm->crumbs_data.mode = LONG;
		kvm->crumbs_data.idt_entry_size=16;
		printk("kvm:start_syscall_trace: system running in long mode (x86_64)\n");
	}

	kvm->crumbs_data.running = 1;

	//code to set #GP trap
	//i=0;
	//while(kvm->vcpus[i] && i<KVM_MAX_VCPUS){
	//kvm_for_each_vcpu(i, vcpu, kvm){
		//vcpu_load(vcpu);
		//kvm_x86_ops->set_gp_trap(vcpu);
		//printk("start_trace_from_host: cpu%d: GP trap set\n",i);
		//vcpu_put(vcpu);

		//i++;
	//}

	//code to cause syscall to cause #UD (64 bit ubuntu)
	//i=0;
	//while(kvm->vcpus[i] && i<KVM_MAX_VCPUS){
	kvm_for_each_vcpu(i, vcpu, kvm){
		vcpu_load(vcpu);
		kvm_get_msr_common(vcpu, MSR_EFER, &(kvm->crumbs_data.efer_val));
		//kvm_get_msr_common(vcpu, MSR_IA32_VMX_MISC, &efer);
		msr.data = kvm->crumbs_data.efer_val & ~EFER_SCE;
		msr.index = MSR_EFER;
		msr.host_initiated = false;
		//kvm_set_msr_common(vcpu, MSR_EFER, kvm->crumbs_data.efer_val & ~EFER_SCE);
		//kvm_set_msr_common(vcpu, MSR_IA32_VMX_MISC, efer);

		kvm_set_msr_common(vcpu, &msr);
		vcpu_put(vcpu);

		//i++;
	}
	strcpy(kvm->crumbs_data.tag, tag);
	printk("Tracing started, tag name is %s\n", kvm->crumbs_data.tag);

	kvm->crumbs_data.ticket = 0;
	kvm->crumbs_data.fdlist_size=0;
	INIT_LIST_HEAD(&(kvm->crumbs_data.fdlisthead.list));
	spin_lock_init(&kvm->crumbs_data.proc_lock);
	//no need to init proc, since kvm module itself will init a global proc
	//if(trace_proc_init(&(kvm->crumbs_data))!=0)
	//	return -1;

	return 0;
}

int stop_trace_from_host(struct kvm * kvm)
{
	struct list_head *pos,*n;
	struct fdnode *p;
	int i;
	//u16 j;
	//struct kvm_sregs sregs;
	//u8 *idt;
	//u64 idt_base,efer;
	//u32 error;
	struct kvm_vcpu *vcpu;
	struct msr_data msr;

	if(!kvm->crumbs_data.running){
		printk("kvm:stop_syscall_trace: WARNING: crumbs is not started, stop will be aborted.\n");
		return 1;
	}

	kvm->crumbs_data.running = 0;

	//code to cause syscall to not cause #UD (64 bit ubuntu)
	//i=0;
	//while(kvm->vcpus[i] && i<KVM_MAX_VCPUS){
	kvm_for_each_vcpu(i, vcpu, kvm){
		vcpu_load(vcpu);
		kvm_get_msr_common(vcpu, MSR_EFER, &(kvm->crumbs_data.efer_val));
		//kvm_get_msr_common(vcpu, MSR_IA32_VMX_MISC, &efer);
		msr.data = kvm->crumbs_data.efer_val | EFER_SCE;
		msr.index = MSR_EFER;
		msr.host_initiated = false;
		//kvm_set_msr_common(vcpu, MSR_EFER, kvm->crumbs_data.efer_val & ~EFER_SCE);
		//kvm_set_msr_common(vcpu, MSR_IA32_VMX_MISC, efer);
	//	rdmsrl_safe(MSR_IA32_VMX_MISC, efer);
	//	printk("start_trace:MSR_IA32_VMX_MISC:%x\n", efer);
	//	printk("start_trace:after:MSR_IA32_VMX_MISC:%x\n", efer|0x20000000);
		kvm_set_msr_common(vcpu, &msr);
		vcpu_put(vcpu);

		//i++;
	}
	kvm->crumbs_data.ticket = 0;
	//empty nodes in fdlist that contains file descriptor
	list_for_each_safe(pos, n, &(kvm->crumbs_data.fdlisthead.list))
	{
		p = list_entry(pos, struct fdnode, list);
		list_del(pos);
		kfree(p);
	}
	if(list_empty(&(kvm->crumbs_data.fdlisthead.list))) {
		printk("fdlist is empty\n");
		kvm->crumbs_data.fdlist_size=0;
		return 0;
	}
	else
		return -1;
	return 0;
}

struct file* file_open(const char* path, int flags, int rights) {
    struct file* filp = NULL;
    mm_segment_t oldfs;
    int err = 0;

    oldfs = get_fs();
    set_fs(get_ds());
    filp = filp_open(path, flags, rights);
    set_fs(oldfs);
    if(IS_ERR(filp)) {
        err = PTR_ERR(filp);
        return NULL;
    }
    return filp;
}
EXPORT_SYMBOL_GPL(file_open);

int file_read(struct file* file, unsigned long long offset, unsigned char* data, unsigned int size) {
    mm_segment_t oldfs;
    int ret;

    oldfs = get_fs();
    set_fs(get_ds());

    ret = vfs_read(file, data, size, &offset);

    set_fs(oldfs);
    return ret;
}
EXPORT_SYMBOL_GPL(file_read);

int file_write(struct file* file, unsigned long long offset, unsigned char* data, unsigned int size) {
    mm_segment_t oldfs;
    int ret;

    oldfs = get_fs();
    set_fs(get_ds());

    ret = vfs_write(file, data, size, &offset);

    set_fs(oldfs);
    return ret;
}
EXPORT_SYMBOL_GPL(file_write);

void file_close(struct file* file) {
    filp_close(file, NULL);
}
EXPORT_SYMBOL_GPL(file_close);

int start_vhost_trace(struct kvm *kvm, int rate)
{
	struct file* filp = NULL;

	/*if(kvm->crumbs_data.running==1){
		printk("kvm:start_vhost_trace: WARNING: crumbs is already running, start will be aborted.\n");
		return 1;
	}*/
	filp = file_open("/proc/vhost-ctl", O_RDWR | O_LARGEFILE, 0);
	if(rate<0)
		rate=0;
	if(rate<10) {
		char str[3];
		snprintf(str, 3, "11%d", rate);
		file_write(filp, 0, str, 3);
	}
	if(rate>=100)
		rate=0;
	if(rate>=10) {
		char str[5];
		snprintf(str, 5, "12%d", rate);
		file_write(filp, 0, str, 5);
	}
	file_close(filp);
	return 0;
}
EXPORT_SYMBOL_GPL(start_vhost_trace);

//disable vhost tracing
int stop_vhost_trace(struct kvm *kvm)
{
	struct file* filp = NULL;

	/*if(!kvm->crumbs_data.running){
		printk("kvm:stop_vhost_trace: WARNING: crumbs is not started, stop will be aborted.\n");
		return 1;
	}*/
	filp = file_open("/proc/vhost-ctl", O_RDWR | O_LARGEFILE, 0);
	file_write(filp, 0, "0", 1);
	file_close(filp);
	return 0;
}
EXPORT_SYMBOL_GPL(stop_vhost_trace);
